open Ast

(* ────────────────────────────────────────────────────────────── *)
(* TAC 数据结构                                                    *)
(* ────────────────────────────────────────────────────────────── *)

type tac =
  | TacAssign of string * string
  | TacBinOp of string * string * string * string
  | TacUnOp of string * string * string
  | TacLabel of string
  | TacGoto of string
  | TacIfGoto of string * string
  | TacParam of string
  | TacCall of string * string * int
  | TacReturn of string option
  | TacComment of string * string * identifier list
  | TacPhi of string * string * string

(* ────────────────────────────────────────────────────────────── *)
(* 全局计数器与栈                                                   *)
(* ────────────────────────────────────────────────────────────── *)

let temp_counter      = ref 0
let if_label_counter  = ref 0
let then_label_counter = ref 0
let while_label_counter = ref 0
let break_stack       = ref []
let continue_stack    = ref []
let current_func      = ref ""

(** 作用域栈：最内层作用域在 head *)
let scope_stack : ((string, string) Hashtbl.t) list ref = ref []

(* ────────────────────────────────────────────────────────────── *)
(* 工具函数                                                         *)
(* ────────────────────────────────────────────────────────────── *)

let new_temp () =
  let t = Printf.sprintf "t%d" !temp_counter in
  incr temp_counter; t

let if_new_label () =
  let l = Printf.sprintf "if_L%d" !if_label_counter in
  incr if_label_counter; l

let then_new_label () =
  let l = Printf.sprintf "then_L%d" !then_label_counter in
  incr then_label_counter; l

let while_new_label () =
  let l = Printf.sprintf "while_L%d" !while_label_counter in
  incr while_label_counter; l

let string_of_binop = function
  | Add -> "+" | Sub -> "-" | Mul -> "*" | Div -> "/" | Mod -> "%"
  | Equal -> "==" | NotEqual -> "!=" | Less -> "<" | LessEqual -> "<="
  | Greater -> ">" | GreaterEqual -> ">="
  | LogicalAnd -> "&&" | LogicalOr -> "||"

let string_of_unop = function
  | UnaryPlus -> "+" | UnaryMinus -> "-" | LogicalNot -> "!"

module SSAMap = Hashtbl
let ssa_version () = Hashtbl.create 32

let ssa_var_name name ver = Printf.sprintf "%s_%d_%s" name ver !current_func

(* ────────────────────────────────────────────────────────────── *)
(* 作用域 / SSA 辅助                                                *)
(* ────────────────────────────────────────────────────────────── *)

(** 把变量插入当前最内层作用域，并返回 SSA 名 *)
let ensure_var_in_current_scope env name =
  match !scope_stack with
  | cur :: _ -> (
      match Hashtbl.find_opt cur name with
      | Some ssa -> ssa
      | None ->
          let v = try SSAMap.find env name with Not_found -> 0 in
          SSAMap.replace env name (v + 1);
          let ssa = ssa_var_name name (v + 1) in
          Hashtbl.add cur name ssa;
          ssa )
  | [] -> failwith "No active scope"

(** 递归从内到外查找某变量的 SSA *)
let rec lookup_var name = function
  | [] -> None
  | scope :: rest -> (
      match Hashtbl.find_opt scope name with
      | Some v -> Some v
      | None -> lookup_var name rest )

(** 根据可见性获取 SSA 名（只读，不自动递增） *)
let get_visible_ssa env name =
  match lookup_var name !scope_stack with
  | Some ssa -> ssa
  | None -> (
      try let v = SSAMap.find env name in ssa_var_name name v
      with Not_found -> ssa_var_name name 0)

(* ────────────────────────────────────────────────────────────── *)
(* 作用域 push / pop                                               *)
(* ────────────────────────────────────────────────────────────── *)

let push_scope () = scope_stack := (Hashtbl.create 32) :: !scope_stack

let pop_scope () =
  match !scope_stack with
  | _ :: rest -> scope_stack := rest
  | [] -> failwith "No scope to pop"

(* ────────────────────────────────────────────────────────────── *)
(* 表达式生成                                                      *)
(* ────────────────────────────────────────────────────────────── *)

let rec gen_expr (e : expression) (env : (string, int) Hashtbl.t) (code : tac list ref) : string =
  match e with
  | Identifier id -> get_visible_ssa env id
  | Number n ->
      let t = new_temp () in
      code := !code @ [TacAssign (t, string_of_int n)];
      t
  | UnaryOp (op, e1) ->
      let t1 = gen_expr e1 env code in
      let t  = new_temp () in
      code := !code @ [TacUnOp (t, string_of_unop op, t1)];
      t
  | BinaryOp (e1, op, e2) ->
      let t1 = gen_expr e1 env code in
      let t2 = gen_expr e2 env code in
      let t  = new_temp () in
      code := !code @ [TacBinOp (t, t1, string_of_binop op, t2)];
      t
  | FunctionCall (fname, args) ->
      let arg_ts = List.map (fun a -> gen_expr a env code) args in
      List.iter (fun t -> code := !code @ [TacParam t]) (List.rev arg_ts);
      let t = new_temp () in
      code := !code @ [TacCall (t, fname, List.length args)];
      t

(* ────────────────────────────────────────────────────────────── *)
(* 语句生成                                                      *)
(* ────────────────────────────────────────────────────────────── *)

let rec gen_stmt (s : statement) (env : (string, int) Hashtbl.t) (code : tac list ref) : unit =
  match s with
  | Block stmts ->
      push_scope ();
      let old_env = Hashtbl.copy env in
      List.iter (fun st -> gen_stmt st env code) stmts;
      Hashtbl.clear env;
      Hashtbl.iter (fun k v -> Hashtbl.add env k v) old_env;
      pop_scope ()
  | EmptyStmt -> ()
  | ExprStmt e -> ignore (gen_expr e env code)
  | Assignment (id, e) ->
      let rhs = gen_expr e env code in
      let ssa_lhs = ensure_var_in_current_scope env id in
      code := !code @ [TacAssign (ssa_lhs, rhs)]
  | VarDecl (id, e) ->
      let rhs = gen_expr e env code in
      let ssa = ensure_var_in_current_scope env id in
      code := !code @ [TacAssign (ssa, rhs)]
  | IfStmt (cond, then_s, else_s_opt) ->
      let cond_t = gen_expr cond env code in
      let cond_not = new_temp () in
      code := !code @ [TacUnOp (cond_not, "!", cond_t)];
      (match else_s_opt with
       | Some else_s ->
           let l_else = if_new_label () in
           let l_end  = if_new_label () in
           code := !code @ [TacIfGoto (cond_not, l_else)];
           let l_then = then_new_label () in
           code := !code @ [TacLabel l_then];
           let env_then = Hashtbl.copy env in
           let code_then = ref [] in
           gen_stmt then_s env_then code_then;
           code_then := !code_then @ [TacGoto l_end];
           let env_else = Hashtbl.copy env in
           let code_else = ref [] in
           code_else := !code_else @ [TacLabel l_else];
           gen_stmt else_s env_else code_else;
           code := !code @ !code_then @ !code_else;
           code := !code @ [TacLabel l_end];
           (* φ 合并 *)
           let collect tbl acc = Hashtbl.fold (fun k _ a -> if List.mem k a then a else k :: a) tbl acc in
           let vars = [] |> collect env_then |> collect env_else in
           let get_ver tbl n = match Hashtbl.find_opt tbl n with Some v -> v | None -> 0 in
           List.iter (fun v_name ->
             let v_then = get_ver env_then v_name in
             let v_else = get_ver env_else v_name in
             if v_then <> v_else then (
               let merged = max v_then v_else + 1 in
               Hashtbl.replace env v_name merged;
               let phi = ssa_var_name v_name merged in
               let then_ssa = ssa_var_name v_name v_then in
               let else_ssa = ssa_var_name v_name v_else in
               code := !code @ [TacPhi (phi, then_ssa, else_ssa)])) vars
       | None ->
           let l_end = if_new_label () in
           code := !code @ [TacIfGoto (cond_not, l_end)];
           gen_stmt then_s env code;
           code := !code @ [TacLabel l_end])
  | WhileStmt (cond, body) ->
      let l_cond = while_new_label () in
      let l_body = while_new_label () in
      let l_end  = while_new_label () in
      code := !code @ [TacGoto l_cond; TacLabel l_cond];
      let cond_t = gen_expr cond env code in
      let cond_not = new_temp () in
      code := !code @ [TacUnOp (cond_not, "!", cond_t); TacIfGoto (cond_not, l_end); TacLabel l_body];
      break_stack    := l_end    :: !break_stack;
      continue_stack := l_cond   :: !continue_stack;
      gen_stmt body env code;
      code := !code @ [TacGoto l_cond; TacLabel l_end];
      break_stack    := List.tl !break_stack;
      continue_stack := List.tl !continue_stack
  | BreakStmt ->
      (match !break_stack with
       | target :: _ -> code := !code @ [TacGoto target]
       | [] -> code := !code @ [TacComment ("a0", "break (not in loop)", [])])
  | ContinueStmt ->
      (match !continue_stack with
       | target :: _ -> code := !code @ [TacGoto target]
       | [] -> code := !code @ [TacComment ("a0", "continue (not in loop)", [])])
  | ReturnStmt eo ->
      (match eo with
       | None -> code := !code @ [TacReturn None]
       | Some e ->
           let rv = gen_expr e env code in
           code := !code @ [TacReturn (Some rv)])

(* ────────────────────────────────────────────────────────────── *)
(* 函数 & 编译单元                                                  *)
(* ────────────────────────────────────────────────────────────── *)

let gen_func (FuncDef (ret_ty, name, params, body)) : tac list =
  current_func := name;
  let param_names = List.map (function Param id -> id) params in
  let tmp = new_temp () in
  let code = ref [TacComment (tmp, "function " ^ name, param_names); TacLabel name] in
  let env = ssa_version () in
  gen_stmt body env code;
  (match ret_ty, List.rev !code with
   | Void, last :: _ -> (
       match last with TacReturn None -> () | _ -> code := !code @ [TacReturn None])
   | _ -> ());
  !code

let gen_comp_unit (cu : comp_unit) : tac list =
  cu |> List.map gen_func |> List.flatten

